'''
作者: Zion Loh
时间: 2025年7月6日
许可: MIT
使用该源代码请务必携带作者版权信息
'''


import sys
import os
import shutil
from collections import Counter, defaultdict
from scapy.all import rdpcap, IP, ARP, DNS, DNSQR, DNSRR, TCP, UDP
import matplotlib.pyplot as plt
from matplotlib import rcParams
import matplotlib
from datetime import datetime
import base64
from io import BytesIO
import tkinter as tk
from tkinter import filedialog, messagebox
import threading
import re

matplotlib.use('Agg') # 使用非交互式后端

rcParams['font.sans-serif'] = ['SimHei']
rcParams['axes.unicode_minus'] = False

def check_tshark_installed():
    tshark = shutil.which("tshark")
    wireshark = shutil.which("wireshark")
    if not tshark or not wireshark:
        messagebox.showerror("环境错误", "未检测到 Wireshark 或 TShark。\n请先安装后再运行程序。")
        sys.exit(1)

def plot_bar_chart(data_dict, title, xlabel, ylabel):
    if not data_dict:
        return ""
    labels, values = zip(*data_dict.items())
    plt.figure(figsize=(12, 6))  # 增大图表尺寸以适应长标签
    plt.bar(labels, values, color='steelblue')
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.xticks(rotation=45, ha='right')  # 自动旋转标签
    plt.tight_layout()  # 自动调整布局，防止标签重叠
    buffer = BytesIO()
    plt.savefig(buffer, format='png')
    plt.close()
    buffer.seek(0)
    img_base64 = base64.b64encode(buffer.read()).decode('utf-8')
    return f'<h3>{title}</h3><img src="data:image/png;base64,{img_base64}"/>'

def plot_protocol_trend(proto_time_dict):
    if not proto_time_dict:
        return ""
    plt.figure(figsize=(14, 7))  # 增大图表尺寸
    for proto, time_series in proto_time_dict.items():
        times = sorted(time_series)
        counts = [time_series[t] for t in times]
        plt.plot(times, counts, label=proto)
    plt.title("协议随时间趋势图")
    plt.xlabel("时间")
    plt.ylabel("报文数量")
    plt.legend()
    plt.xticks(rotation=45)  # 增加 x 轴的旋转角度以避免标签重叠
    plt.tight_layout()  # 自动调整布局，防止标签重叠
    buffer = BytesIO()
    plt.savefig(buffer, format='png')
    plt.close()
    buffer.seek(0)
    img_base64 = base64.b64encode(buffer.read()).decode('utf-8')
    return f'<h3>协议趋势图</h3><img src="data:image/png;base64,{img_base64}"/>'


def analyze_pcap(file_path, output_label):
    if not os.path.exists(file_path):
        output_label.config(text=f"[错误] 文件不存在: {file_path}")
        return

    output_label.config(text=f"正在分析: {file_path} ...")
    try:
        packets = rdpcap(file_path)
    except Exception as e:
        output_label.config(text=f"[错误] 读取文件失败: {e}")
        return

    total_packets = len(packets)
    ip_packets = [pkt for pkt in packets if IP in pkt]

    protocol_counts = Counter()
    src_ip_counts = Counter()
    dst_ip_counts = Counter()
    broadcast_count = 0
    multicast_count = 0
    suspicious_packets = []
    proto_time_dict = defaultdict(lambda: defaultdict(int))
    abnormal_ips = set()

    dns_queries = Counter()
    dns_responses = Counter()
    http_hosts = Counter()
    https_count = 0
    http_methods = Counter()

    for pkt in packets:
        try:
            timestamp = datetime.fromtimestamp(float(pkt.time)).strftime('%H:%M')
        except Exception:
            timestamp = "Unknown"

        if IP in pkt:
            proto = pkt[IP].proto
            src = pkt[IP].src
            dst = pkt[IP].dst

            if proto == 1:
                proto_name = 'ICMP'
            elif proto == 6:
                proto_name = 'TCP'
            elif proto == 17:
                proto_name = 'UDP'
            else:
                proto_name = f'Other({proto})'

            protocol_counts[proto_name] += 1
            src_ip_counts[src] += 1
            dst_ip_counts[dst] += 1
            proto_time_dict[proto_name][timestamp] += 1

            if dst.startswith("224.") or dst.startswith("239.") or dst.startswith("ff"):
                multicast_count += 1
            elif dst.endswith(".255") or dst == "255.255.255.255":
                broadcast_count += 1

            if proto_name.startswith("Other"):
                suspicious_packets.append(pkt.summary())

            if src_ip_counts[src] > 100 or dst_ip_counts[dst] > 100:
                abnormal_ips.add(src)
                abnormal_ips.add(dst)

            if pkt.haslayer(DNS) and pkt.haslayer(UDP):
                if pkt[DNS].qr == 0 and pkt.haslayer(DNSQR):
                    qname = pkt[DNSQR].qname.decode(errors='ignore')
                    dns_queries[qname] += 1
                elif pkt[DNS].qr == 1 and pkt.haslayer(DNSRR):
                    rname = pkt[DNSRR].rrname.decode(errors='ignore')
                    dns_responses[rname] += 1

            if pkt.haslayer(TCP):
                sport = pkt[TCP].sport
                dport = pkt[TCP].dport
                if dport == 443 or sport == 443:
                    https_count += 1
                if pkt.haslayer('Raw'):
                    try:
                        raw_data = pkt['Raw'].load.decode('utf-8', errors='ignore')
                        if raw_data.startswith(('GET', 'POST', 'HEAD', 'PUT', 'DELETE')):
                            lines = raw_data.split("\r\n")
                            method = lines[0].split()[0]
                            http_methods[method] += 1
                            for line in lines:
                                if line.lower().startswith('host:'):
                                    host = line.split(":", 1)[1].strip()
                                    http_hosts[host] += 1
                    except Exception:
                        pass

        elif ARP in pkt:
            protocol_counts['ARP'] += 1
            proto_time_dict['ARP'][timestamp] += 1

    html_parts = [
        "<html><head><meta charset='utf-8'><style>body{font-family:sans-serif;padding:2em;}img{max-width:100%;}</style></head><body>",
        f"<h2>抓包分析报告 - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</h2>",
        f"<p><strong>文件:</strong> {os.path.basename(file_path)}</p>",
        f"<p><strong>总报文:</strong> {total_packets}，<strong>IP 报文:</strong> {len(ip_packets)}，<strong>广播:</strong> {broadcast_count}，<strong>多播:</strong> {multicast_count}</p>",
        plot_bar_chart(protocol_counts, "协议分布统计", "协议", "数量"),
        plot_bar_chart(dict(src_ip_counts.most_common(10)), "Top 10 源 IP", "IP", "数量"),
        plot_bar_chart(dict(dst_ip_counts.most_common(10)), "Top 10 目标 IP", "IP", "数量"),
        plot_protocol_trend(proto_time_dict),
    ]

    if suspicious_packets:
        html_parts.append("<h3>异常协议报文（前10条）</h3><ul>")
        for s in suspicious_packets[:10]:
            html_parts.append(f"<li>{s}</li>")
        html_parts.append("</ul>")

    if abnormal_ips:
        abnormal_count = Counter()
        for ip in abnormal_ips:
            abnormal_count[ip] = src_ip_counts[ip] + dst_ip_counts[ip]
        top_abnormals = abnormal_count.most_common(20)
        html_parts.append(f"<h3>异常 IP 地址（Top 20）</h3><ul>")
        for ip, count in top_abnormals:
            html_parts.append(f"<li>{ip}：{count} 次</li>")
        html_parts.append("</ul>")

    if dns_queries:
        html_parts.append(plot_bar_chart(dict(dns_queries.most_common(10)), "Top 10 DNS 查询", "域名", "次数"))
    if dns_responses:
        html_parts.append(plot_bar_chart(dict(dns_responses.most_common(10)), "Top 10 DNS 响应", "域名", "次数"))

    if http_hosts:
        html_parts.append(plot_bar_chart(dict(http_hosts.most_common(10)), "Top 10 HTTP Host", "Host", "次数"))
    if http_methods:
        html_parts.append(plot_bar_chart(http_methods, "HTTP 方法分布", "方法", "数量"))
    if https_count:
        html_parts.append(f"<p><strong>HTTPS 会话数量:</strong> {https_count}</p>")

    html_parts.append("<h3>分析建议</h3><ul>")
    if broadcast_count > 100:
        html_parts.append("<li>广播报文较多，建议检查是否存在广播风暴或DHCP泛洪。</li>")
    if multicast_count > 100:
        html_parts.append("<li>多播报文较多，建议启用IGMP Snooping防止泛洪。</li>")
    if suspicious_packets:
        html_parts.append("<li>存在异常协议（非常用TCP/UDP/ICMP），请进一步确认其合法性。</li>")
    if sum(dns_queries.values()) > 100:
        html_parts.append("<li>DNS 查询数量异常偏高，可能存在恶意域名投毒或隧道通信。</li>")
    if https_count > 1000:
        html_parts.append("<li>HTTPS 会话数量较多，请确认是否有加密流量攻击或爬虫行为。</li>")
    if not suspicious_packets and broadcast_count < 50 and multicast_count < 50:
        html_parts.append("<li>网络通信整体正常，未发现明显异常。</li>")
    html_parts.append("</ul></body></html>")

    report_path = os.path.splitext(file_path)[0] + "_report.html"
    with open(report_path, "w", encoding="utf-8") as f:
        f.write("".join(html_parts))
    output_label.config(text=f"[完成] 报告已生成：{report_path}")

    # 弹窗提示是否打开报告
    if messagebox.askyesno("报告生成完成", "报告已生成，是否立即打开？"):
        if sys.platform == "win32":
            os.startfile(report_path)  # Windows
        elif sys.platform == "darwin":
            subprocess.call(["open", report_path])  # macOS
        else:
            webbrowser.open(f'file://{os.path.realpath(report_path)}')  # Linux

def select_file(output_label):
    file_path = filedialog.askopenfilename(filetypes=[
        ("Wireshark Files", "*.pcap *.pcapng"),
        ("PCAP", "*.pcap"),
        ("PCAPNG", "*.pcapng"),
        ("All Files", "*.*")
    ])
    if file_path:
        threading.Thread(target=analyze_pcap, args=(file_path, output_label), daemon=True).start()

def create_gui():
    root = tk.Tk()
    root.title("抓包文件分析工具 作者:Zion Loh https://gitee.com/zionloh")
    width = 600
    height = 220
    screen_width = root.winfo_screenwidth()
    screen_height = root.winfo_screenheight()
    position_top = int(screen_height / 2 - height / 2)
    position_right = int(screen_width / 2 - width / 2)
    root.geometry(f'{width}x{height}+{position_right}+{position_top}')
    root.resizable(True, True)

    label = tk.Label(root, text="请选择一个抓包文件进行分析", font=("Arial", 14), wraplength=700, justify="center")
    label.pack(pady=20)
    output_label = tk.Label(root, text="", font=("Arial", 12), wraplength=500, justify="center")
    output_label.pack(pady=20)
    select_button = tk.Button(root, text="选择文件", command=lambda: select_file(output_label), font=("Arial", 12))
    select_button.pack(pady=10)

    root.mainloop()

if __name__ == "__main__":
    create_gui()
